# KAN用于多参数估计
from typing import Dict, Callable
from tqdm import tqdm
import matplotlib
import matplotlib.pyplot as plt
import torch
from torch.utils.data import DataLoader
import sympy as sp
from sympy.abc import x
from kan.utils import create_dataset
from kan.utils import ex_round
from kan import KAN
from apps.fmcw.conf.app_config import AppConfig as AF
from dss.fmcw_mpe_dss import FmcwMpeDss

class MpeKan(object):
    def __init__(self):
        self.name = 'apps.fmcw.ias.mpe_kan.MpeKan'
        # 定义网络结构
        # 指定字体名称
        matplotlib.rcParams['font.family'] = 'SimHei'
        matplotlib.rcParams['axes.unicode_minus'] = False  # 正确显示负号
        torch.set_default_dtype(torch.float64)
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        kan_ckpts_path = './work/mpe_kan/ckpts'
        kan_figures_path = './work/mpe_kan/images'
        L0 = AF.numADC * AF.numChirps * AF.numTX * AF.numRX * 2 # 将复数变为实数
        L1 = L0 // (8192*32) # 16384
        L2 = 12 # L1 // 8 # 2048
        print(f'模型大小：{(L0*L1*L2)/1024/1024/1024}G ??????')
        L3 = L2 // 4 # 512
        L4 = 12 # 三个目标，每个目标４个参数：概率、径向距离、径向速度、水平到达角
        print(f'开始创建KAN2网络')
        self.model = KAN(width=[L0, L1, L2], grid=3, k=3, seed=42, device=device, ckpt_path=kan_ckpts_path)
        self.model = self.model.to(device)
        print(f'创建KAN2模型')
        training_data = FmcwMpeDss()
        test_data = FmcwMpeDss()
        train_dataloader = DataLoader(training_data, batch_size=1, shuffle=True)
        test_dataloader = DataLoader(test_data, batch_size=1, shuffle=True)
        X, y = next(iter(train_dataloader))
        print(f'X: {X.shape}; y: {y.shape};')
        epochs = 100
        pbar = tqdm(range(epochs), desc='description', ncols=100)
        for epoch in pbar:
            self.train(dataloader=train_dataloader, model=self.model, optimizer='LBFGS')
        print(f'^_^ The End! ^_^')

        # # create dataset f(x,y) = exp(sin(pi*x)+y^2)
        # f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]]) + x[:,[1]]**2)
        # dataset = create_dataset(f, n_var=2, device=device) # Dict(tensor): train_input, test_input, train_label, test_label

        # #dataset['train_input'].shape, dataset['train_label'].shape
        # print(f'dataset: ')
        # print(f'    train_input: {dataset["train_input"].shape};')
        # print(f'    train_label: {dataset["train_label"].shape};')
        # print(f'    test_input: {dataset["test_input"].shape};')
        # print(f'    test_label: {dataset["test_label"].shape};')

    def train(self, dataloader, model, loss_fn:Callable = None, optimizer:str = 'LBFGS', device:str = 'cuda') -> None:
        model.train()
        for bidx, (X, y) in enumerate(dataloader):
            X, y = X.to(device), y.to(device)
            dataset = {
                'train_input': X,
                'train_label': y,
                'test_input': X,
                'test_label': y
            }
            model(dataset['train_input'])
            model.fit(dataset, opt=optimizer, steps=1, lamb=0.001)

    def train_step(self) -> None:
        pass

    def test(self) -> None:
        pass

    @staticmethod
    def learn(params:Dict = {}) -> None:
        # 复现KAN的ＨelloＷorld
        print(f'Hello KAN2 v0.0.1')
        kan_ckpts_path = './work/mpe_kan/ckpts'
        kan_figures_path = './work/mpe_kan/images'
        # 指定字体名称
        matplotlib.rcParams['font.family'] = 'SimHei'
        matplotlib.rcParams['axes.unicode_minus'] = False  # 正确显示负号
        torch.set_default_dtype(torch.float64)
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(device)
        # create a KAN: 2D inputs, 1D output, and 5 hidden neurons. cubic spline (k=3), 5 grid intervals (grid=3).
        model = KAN(width=[2,5,1], grid=3, k=3, seed=42, device=device, ckpt_path=kan_ckpts_path)
        # create dataset f(x,y) = exp(sin(pi*x)+y^2)
        f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]]) + x[:,[1]]**2)
        dataset = create_dataset(f, n_var=2, device=device)
        dataset['train_input'].shape, dataset['train_label'].shape
        # plot KAN at initialization
        model(dataset['train_input'])
        model.plot(folder=kan_figures_path)
        plt.show()
        # train the model
        model.fit(dataset, opt="LBFGS", steps=50, lamb=0.001)
        model.plot(folder=kan_figures_path)
        plt.show()
        # Prune and replot
        model = model.prune()
        model.plot(folder=kan_figures_path)
        plt.show()
        # Continue training and replot
        model.fit(dataset, opt="LBFGS", steps=50)
        model = model.refine(10)
        model.fit(dataset, opt="LBFGS", steps=50)
        model.plot(folder=kan_figures_path)
        plt.show()
        # Automatically or manually set activation functions to be symbolic
        mode = "auto" # "manual"
        if mode == "manual":
            # manual mode
            model.fix_symbolic(0,0,0,'sin');
            model.fix_symbolic(0,1,0,'x^2');
            model.fix_symbolic(1,0,0,'exp');
        elif mode == "auto":
            # automatic mode
            lib = ['x','x^2','x^3','x^4','exp','log','sqrt','tanh','sin','abs']
            model.auto_symbolic(lib=lib)
        model.fit(dataset, opt="LBFGS", steps=50)
        print(ex_round(model.symbolic_formula()[0][0],4))